from typing import TypedDict, List
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langchain.schema import HumanMessage
from dotenv import load_dotenv
import os

from langchain.chat_models import init_chat_model

# 初始化LLM（请确保已安装对应依赖并配置API密钥）
llm = init_chat_model(model="deepseek-chat", model_provider="deepseek")


# 1. 定义多智能体状态
class AgentState(TypedDict):
    research_materials: str  # 研究员收集的资料
    draft: str  # 写作助手的初稿
    review_result: str  # 审核员的审核结果
    is_draft_approved: bool  # 初稿是否通过审核
    task: str  # 核心任务（固定为“生成LangGraph教程”）
    iteration_count: int  # 迭代次数计数器，防止无限循环


# 2. 定义智能体节点
# 研究员节点：收集资料（根据审核结果补充）
def researcher_agent(state: AgentState) -> AgentState:
    task = state["task"]
    review_result = state["review_result"]
    existing_materials = state["research_materials"]
    iteration = state["iteration_count"]

    # 根据审核结果调整收集策略
    if review_result and not state["is_draft_approved"]:
        research_prompt = f"""任务：{task}。
        之前的资料被审核发现问题：{review_result}。
        请基于已有资料补充缺失内容（无需重复已有信息），重点完善上述问题。
        已有资料：{existing_materials}
        注意：这是第{iteration}次补充，请确保覆盖所有缺失点。
        """
    else:
        research_prompt = f"""任务：{task}。
        请收集LangGraph的核心知识点，必须包含：
        1. 定义（是什么）
        2. 核心组件（如节点、边、状态等具体说明）
        3. 使用步骤（从创建到运行的完整流程）
        确保信息全面，不遗漏关键部分。
        """

    research_result = llm.invoke([HumanMessage(content=research_prompt)]).content
    # 合并新旧资料（首次收集直接使用新资料）
    updated_materials = f"{existing_materials}\n补充资料：{research_result}" if existing_materials else research_result
    return {
        "research_materials": updated_materials,
        "iteration_count": iteration + 1  # 迭代次数+1
    }


# 写作助手节点：撰写初稿
def writer_agent(state: AgentState) -> AgentState:
    materials = state["research_materials"]
    write_prompt = f"基于以下资料撰写LangGraph教程（300字左右）：\n{materials}\n要求：结构清晰，语言通俗，必须包含定义、核心组件和使用步骤三部分。"
    draft = llm.invoke([HumanMessage(content=write_prompt)]).content
    return {"draft": draft}


# 审核员节点：审核初稿（明确标准）
def reviewer_agent(state: AgentState) -> AgentState:
    draft = state["draft"]
    review_prompt = f"""严格审核以下LangGraph教程初稿是否完整，必须包含3个核心部分：
    1. 定义（明确说明LangGraph是什么）
    2. 核心组件（列出节点、边、状态等具体组件并简要说明）
    3. 使用步骤（从初始化到运行的完整流程，至少3个步骤）

    审核结果格式要求：
    - 若完整：仅返回"通过"
    - 若不完整：仅返回缺失的部分（例如："缺失核心组件的具体说明"）

    待审核初稿：{draft}
    """
    review_response = llm.invoke([HumanMessage(content=review_prompt)]).content
    is_approved = "通过" in review_response
    return {
        "review_result": review_response,
        "is_draft_approved": is_approved
    }


# 3. 构建多智能体图
graph = StateGraph(AgentState)

# 添加节点
graph.add_node("researcher", researcher_agent)  # 研究员
graph.add_node("writer", writer_agent)  # 写作助手
graph.add_node("reviewer", reviewer_agent)  # 审核员

# 定义边：协作流程
graph.set_entry_point("researcher")  # 入口：先启动研究员
graph.add_edge("researcher", "writer")  # 研究员→写作助手（资料收集后撰写）
graph.add_edge("writer", "reviewer")  # 写作助手→审核员（初稿完成后审核）

# 审核后的条件分支：通过则结束，否则返回研究员补充资料
graph.add_conditional_edges(
    "reviewer",
    lambda state: "researcher" if not state["is_draft_approved"] else END
)

# 4. 运行多智能体系统（增加递归限制和迭代计数）
compiled_graph = graph.compile()  # 限制最大递归次数为10次
initial_state = {
    "research_materials": "",
    "draft": "",
    "review_result": "",
    "is_draft_approved": False,
    "task": "生成LangGraph教程（面向新手，300字左右）",
    "iteration_count": 1  # 初始迭代次数为1
}

# 执行流程
result = compiled_graph.invoke(initial_state)

# 打印最终结果
print("=== 多智能体协作最终结果 ===")
print(f"1. 研究员资料：\n{result['research_materials']}\n")
print(f"2. 写作初稿：\n{result['draft']}\n")
print(f"3. 审核结果：\n{result['review_result']}\n")
print(f"4. 总迭代次数：{result['iteration_count'] - 1}")
